import torch
import torch.nn as nn
import torch.optim as optim
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from torch.utils.data import DataLoader, TensorDataset
import matplotlib.pyplot as plt
import numpy as np
#导入模块

dataset = load_iris()
x = torch.tensor(dataset['data'], dtype=torch.float32)
y = torch.tensor(dataset['target'], dtype=torch.long)
x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=0.4)
train_dataset = TensorDataset(x_train, y_train)
#加载数据集，以及相应的输出

class IRIS(nn.Module):
    def __init__(self):
        super(IRIS, self).__init__()
        self.fc1 = nn.Linear(4, 16)
        self.fc2 = nn.Linear(16, 3)
    def forward(self,x):
        x = torch.tanh(self.fc1(x))
        x = self.fc2(x)
        return x
model = IRIS()
#使用torch框架中的神经网络模型，来构建框架

lost = []
epochs = []
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adamax(model.parameters(), lr=0.08)
num_epochs = 100
for epoch in range(num_epochs):
    optimizer.zero_grad()
    output = model(x_train)
    loss = criterion(output, y_train)
    loss.backward()
    optimizer.step()
    if (epoch+1) % 10 == 0:
        print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item()}')
    epochs.append(epoch)
    lost.append(loss.item())
y = np.array(lost)
x = np.array(epochs)
#模型的迭代以及为后续画图做准备

with torch.no_grad():
    model.eval()
    y_pred = model(x_test)
    _, predicted = torch.max(y_pred, 1)
correct = (predicted == y_test).sum().item()
total = y_test.size(0)
accuracy = correct / total * 100
print(f'Accuracy on test set: {accuracy:.4f}%')
torch.save(model, 'iris.pt')
#对模型进行评估以及保存模型

plt.xlabel('Epoch')
plt.ylabel('loss')
plt.plot(x, y)
plt.show()
#画出训练集随着迭代次数的增加，loss的值的变化